"""
简化的黄金交易预测特征工程脚本
"""

import pandas as pd
import numpy as np
import ta
from datetime import datetime
import os


def create_technical_indicators(data):
    """创建技术指标"""
    print("创建技术指标...")
    
    df = data.copy()
    df = df.sort_values('date').reset_index(drop=True)
    
    # 移动平均线
    df['sma_5'] = ta.trend.sma_indicator(df['close'], window=5)
    df['sma_10'] = ta.trend.sma_indicator(df['close'], window=10)
    df['sma_20'] = ta.trend.sma_indicator(df['close'], window=20)
    df['sma_50'] = ta.trend.sma_indicator(df['close'], window=50)
    
    # 指数移动平均线
    df['ema_5'] = ta.trend.ema_indicator(df['close'], window=5)
    df['ema_10'] = ta.trend.ema_indicator(df['close'], window=10)
    df['ema_20'] = ta.trend.ema_indicator(df['close'], window=20)
    
    # MACD
    df['macd'] = ta.trend.macd(df['close'])
    df['macd_signal'] = ta.trend.macd_signal(df['close'])
    df['macd_diff'] = ta.trend.macd_diff(df['close'])
    
    # RSI
    df['rsi_14'] = ta.momentum.rsi(df['close'], window=14)
    
    # 布林带
    df['bb_upper'] = ta.volatility.bollinger_hband(df['close'])
    df['bb_middle'] = ta.volatility.bollinger_mavg(df['close'])
    df['bb_lower'] = ta.volatility.bollinger_lband(df['close'])
    df['bb_width'] = (df['bb_upper'] - df['bb_lower']) / df['bb_middle']
    
    # 随机指标
    df['stoch_k'] = ta.momentum.stoch(df['high'], df['low'], df['close'])
    df['stoch_d'] = ta.momentum.stoch_signal(df['high'], df['low'], df['close'])
    
    # 威廉指标
    df['williams_r'] = ta.momentum.williams_r(df['high'], df['low'], df['close'])
    
    # 商品通道指数
    df['cci'] = ta.trend.cci(df['high'], df['low'], df['close'])
    
    # 平均真实波幅
    df['atr'] = ta.volatility.average_true_range(df['high'], df['low'], df['close'])
    
    # 动量指标
    df['momentum_1'] = ta.momentum.roc(df['close'], window=1)
    df['momentum_5'] = ta.momentum.roc(df['close'], window=5)
    
    # 价格变化率
    df['price_change_1'] = df['close'].pct_change(1)
    df['price_change_5'] = df['close'].pct_change(5)
    df['price_change_10'] = df['close'].pct_change(10)
    
    # 对数收益率
    df['log_return_1'] = np.log(df['close'] / df['close'].shift(1))
    df['log_return_5'] = np.log(df['close'] / df['close'].shift(5))
    
    # 波动率
    df['volatility_5'] = df['log_return_1'].rolling(window=5).std()
    df['volatility_10'] = df['log_return_1'].rolling(window=10).std()
    
    print(f"创建了 {len(df.columns)} 个特征")
    return df


def create_time_features(data):
    """创建时间特征"""
    print("创建时间特征...")
    
    df = data.copy()
    df['date'] = pd.to_datetime(df['date'])
    
    # 基础时间特征
    df['year'] = df['date'].dt.year
    df['month'] = df['date'].dt.month
    df['day'] = df['date'].dt.day
    df['dayofweek'] = df['date'].dt.dayofweek
    df['quarter'] = df['date'].dt.quarter
    
    # 周期性特征
    df['month_sin'] = np.sin(2 * np.pi * df['month'] / 12)
    df['month_cos'] = np.cos(2 * np.pi * df['month'] / 12)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df['dayofweek'] / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df['dayofweek'] / 7)
    
    # 是否为月初/月末
    df['is_month_start'] = df['date'].dt.is_month_start.astype(int)
    df['is_month_end'] = df['date'].dt.is_month_end.astype(int)
    df['is_weekend'] = (df['dayofweek'] >= 5).astype(int)
    
    print(f"创建了时间特征")
    return df


def create_lag_features(data, target_col='close', lags=[1, 2, 3, 5, 10]):
    """创建滞后特征"""
    print("创建滞后特征...")
    
    df = data.copy()
    
    for lag in lags:
        df[f'{target_col}_lag_{lag}'] = df[target_col].shift(lag)
    
    print(f"创建了 {len(lags)} 个滞后特征")
    return df


def create_rolling_features(data, target_col='close', windows=[5, 10, 20]):
    """创建滚动窗口特征"""
    print("创建滚动窗口特征...")
    
    df = data.copy()
    
    for window in windows:
        # 滚动统计
        df[f'{target_col}_rolling_mean_{window}'] = df[target_col].rolling(window).mean()
        df[f'{target_col}_rolling_std_{window}'] = df[target_col].rolling(window).std()
        df[f'{target_col}_rolling_min_{window}'] = df[target_col].rolling(window).min()
        df[f'{target_col}_rolling_max_{window}'] = df[target_col].rolling(window).max()
        
        # 相对位置
        df[f'{target_col}_rolling_position_{window}'] = (df[target_col] - df[f'{target_col}_rolling_min_{window}']) / (df[f'{target_col}_rolling_max_{window}'] - df[f'{target_col}_rolling_min_{window}'])
    
    print(f"创建了 {len(windows) * 5} 个滚动特征")
    return df


def create_target_variables(data, target_col='close', forecast_horizons=[1, 5, 10]):
    """创建目标变量"""
    print("创建目标变量...")
    
    df = data.copy()
    
    for horizon in forecast_horizons:
        # 未来收益率
        df[f'target_return_{horizon}d'] = df[target_col].shift(-horizon) / df[target_col] - 1
        
        # 未来价格方向
        df[f'target_direction_{horizon}d'] = (df[f'target_return_{horizon}d'] > 0).astype(int)
        
        # 未来价格是否超过阈值
        threshold = 0.02  # 2%的阈值
        df[f'target_up_{horizon}d'] = (df[f'target_return_{horizon}d'] > threshold).astype(int)
        df[f'target_down_{horizon}d'] = (df[f'target_return_{horizon}d'] < -threshold).astype(int)
    
    print(f"创建了 {len(forecast_horizons) * 4} 个目标变量")
    return df


def process_combined_data(data):
    """处理合并的黄金和美元指数数据"""
    print("处理合并数据...")
    
    df = data.copy()
    
    # 为黄金数据创建特征
    gold_data = pd.DataFrame({
        'date': df['date'],
        'open': df['gold_close'],
        'high': df['gold_close'] * 1.01,
        'low': df['gold_close'] * 0.99,
        'close': df['gold_close'],
        'volume': np.random.randint(100000, 500000, len(df))
    })
    
    # 为美元指数创建特征
    usd_data = pd.DataFrame({
        'date': df['date'],
        'open': df['usd_close'],
        'high': df['usd_close'] * 1.005,
        'low': df['usd_close'] * 0.995,
        'close': df['usd_close'],
        'volume': np.random.randint(50000, 200000, len(df))
    })
    
    # 创建黄金特征
    gold_features = create_technical_indicators(gold_data)
    gold_features = create_time_features(gold_features)
    gold_features = create_lag_features(gold_features, 'close')
    gold_features = create_rolling_features(gold_features, 'close')
    
    # 创建美元指数特征
    usd_features = create_technical_indicators(usd_data)
    usd_features = create_time_features(usd_features)
    
    # 重命名美元指数特征列
    usd_feature_cols = [col for col in usd_features.columns if col not in ['date', 'open', 'high', 'low', 'close', 'volume']]
    usd_features = usd_features[['date'] + usd_feature_cols]
    usd_features = usd_features.rename(columns={col: f'usd_{col}' for col in usd_feature_cols})
    
    # 合并所有特征
    final_features = pd.merge(gold_features, usd_features, on='date', how='left')
    
    # 创建目标变量
    final_features = create_target_variables(final_features, 'close')
    
    print(f"最终特征数据形状: {final_features.shape}")
    return final_features


def save_features(data, filename, directory="data/features"):
    """保存特征数据"""
    os.makedirs(directory, exist_ok=True)
    filepath = os.path.join(directory, filename)
    data.to_csv(filepath, index=False, encoding='utf-8-sig')
    print(f"特征数据已保存到: {filepath}")
    return filepath


def main():
    """主函数"""
    print("=" * 60)
    print("黄金交易预测特征工程程序（简化版）")
    print("=" * 60)
    
    # 加载原始数据
    data_dir = "data/raw"
    combined_file = os.path.join(data_dir, "combined_data.csv")
    
    if not os.path.exists(combined_file):
        print(f"数据文件不存在: {combined_file}")
        return
    
    print("加载原始数据...")
    combined_data = pd.read_csv(combined_file)
    combined_data['date'] = pd.to_datetime(combined_data['date'])
    print(f"加载数据: {combined_data.shape}")
    
    # 处理数据
    features = process_combined_data(combined_data)
    
    # 保存特征数据
    save_features(features, 'gold_features.csv')
    
    # 显示特征信息
    print("\n特征数据信息:")
    print(f"数据形状: {features.shape}")
    print(f"特征数量: {len(features.columns)}")
    
    print("\n前5行数据:")
    print(features[['date', 'close']].head())
    
    print("\n特征统计:")
    feature_cols = [col for col in features.columns if col not in ['date', 'open', 'high', 'low', 'close', 'volume']]
    print(f"技术指标特征: {len([col for col in feature_cols if any(x in col for x in ['sma', 'ema', 'macd', 'rsi', 'bb', 'stoch'])])}")
    print(f"时间特征: {len([col for col in feature_cols if any(x in col for x in ['year', 'month', 'day', 'quarter', 'sin', 'cos'])])}")
    print(f"滞后特征: {len([col for col in feature_cols if 'lag' in col])}")
    print(f"滚动特征: {len([col for col in feature_cols if 'rolling' in col])}")
    print(f"目标变量: {len([col for col in feature_cols if 'target' in col])}")
    
    # 显示一些关键特征的统计信息
    print("\n关键特征统计:")
    key_features = ['close', 'sma_20', 'rsi_14', 'macd', 'volatility_10', 'target_return_1d']
    for feature in key_features:
        if feature in features.columns:
            print(f"{feature}: 均值={features[feature].mean():.4f}, 标准差={features[feature].std():.4f}")
    
    print("\n特征工程完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()
